import torch
from torch.distributions import multinomial

import d2l

fair_probs = torch.ones([6]) / 6

total_count = 1000
counts = multinomial.Multinomial(total_count, probs=fair_probs).sample((1000,))
print('counts:')
print(counts)

cum_counts = counts.cumsum(dim=0)
print('cum_counts:')
print(cum_counts)

print('cum_counts.sum(dim=1, keepdims=True):')
print(cum_counts.sum(dim=1, keepdims=True))

estimates = cum_counts / cum_counts.sum(dim=1, keepdims=True)
print('estimates:')
print(estimates)

d2l.set_figsize((6, 4.5))
for i in range(6):
    d2l.plt.plot(estimates[:, i].numpy(), label=f'P(die={i + 1})')
d2l.plt.axhline(y = 0.167, color='black', linestyle='dashed')
d2l.plt.gca().set_xlabel('Groups of experiments')
d2l.plt.gca().set_ylabel('Estimated probability')
d2l.plt.legend()
d2l.plt.show()
